import torch 
import torch.nn as nn 


class Conv2d(nn.Module):
    def __init__(self, nin, nout, ks=1, st=1):
        super().__init__()
        pad = (ks-1)//2
        self.layers = nn.Sequential(
            nn.Conv2d(nin, nout, ks, st, padding=pad), 
            nn.BatchNorm2d(nout), 
            nn.LeakyReLU(), 
        )
    def forward(self, x):
        x = self.layers(x)
        return x 
class ResBlock(nn.Module):
    def __init__(self, nin, nmid):
        super().__init__()
        self.layers = nn.Sequential(
            Conv2d(nin, nmid, ks=1, st=1), 
            Conv2d(nmid, nin, ks=3, st=1), 
        )
    def forward(self, x):
        y = self.layers(x) 
        return x + y 
class RepeatBlock(nn.Module):
    def __init__(self, nin, nrepeat):
        super().__init__()
        self.layers = nn.ModuleList()
        for i in range(nrepeat):
            self.layers.append(
                ResBlock(nin, nin//2)
            )
    def forward(self, x):
        for layer in self.layers:
            x = layer(x) 
        return x  
class OutLayer(nn.Module):
    def __init__(self, nin, nout):
        super().__init__() 
        self.layers = nn.Sequential(
            Conv2d(nin, nout//2, 1, 1), 
            Conv2d(nout//2, nout, 3, 1), 
            Conv2d(nout, nout//2, 1, 1), 
            Conv2d(nout//2, nout, 3, 1), 
            Conv2d(nout, nout//2, 1, 1)
        )
    def forward(self, x):
        x = self.layers(x) 
        return x 
class YoloLayer1(nn.Module):
    def __init__(self, nin, nout, branch=True):
        super().__init__() 
        self.layers = nn.Sequential(
            OutLayer(nin, nout), 
        )
        self.out = nn.Sequential(
            Conv2d(nout//2, nout, 3, 1), 
            nn.Conv2d(nout, 255, kernel_size=1), 
        )
        self.brt = nn.Sequential(
                Conv2d(nout//2, nout//4, 1, 1), 
                nn.UpsamplingNearest2d(scale_factor=2), 
            )
    def forward(self, x):
        h = self.layers(x) 
        y1 = self.out(h)
        y2 = self.brt(h)
        return y1, y2
class YoloLayer2(nn.Module):
    def __init__(self, nin, nout, branch=True):
        super().__init__() 
        self.layers = nn.Sequential(
            OutLayer(nin, nout), 
        )
        self.out = nn.Sequential(
            Conv2d(nout//2, nout, 3, 1), 
            nn.Conv2d(nout, 255, kernel_size=1), 
        )
    def forward(self, x):
        h = self.layers(x) 
        y = self.out(h)
        return y 
def _make_grid(nx=20, ny=20):
    yv, xv = torch.meshgrid(torch.arange(ny), torch.arange(nx))
    return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
def mkoutput(x, stride=8):
    B, C, H, W = x.shape 
    x = x.reshape(B, 3, 85, H, W).permute(0, 1, 3, 4, 2).contiguous()
    return x 
class YoloModel(nn.Module):
    """YOLOv3 object detection model"""
    def __init__(self):
        super().__init__()
        self.base0 = nn.Sequential(
            Conv2d(3, 32, 3, 1), 
            Conv2d(32, 64, 3, 2), 
            ResBlock(64, 32), 
            Conv2d(64, 128, 3, 2), 
            RepeatBlock(128, 2), 
            Conv2d(128, 256, 3, 2), 
            RepeatBlock(256, 3), 
        )
        self.base1 = nn.Sequential(
            Conv2d(256, 512, 3, 2), 
            RepeatBlock(512, 8)
        )
        self.base2 = nn.Sequential(
            Conv2d(512, 1024, 3, 2), 
            RepeatBlock(1024, 4)
        )

        self.yolo2 = YoloLayer1(1024, 1024, True) 
        self.yolo1 = YoloLayer1(256*3, 512, True) 
        self.yolo0 = YoloLayer2(128*3, 256, False)
    def forward(self, x):
        h0 = self.base0(x) 
        h1 = self.base1(h0) 
        h2 = self.base2(h1) 
        
        y2, cat1 = self.yolo2(h2) 
        h1 = torch.cat([h1, cat1], dim=1)
        y1, cat0 = self.yolo1(h1) 
        h0 = torch.cat([h0, cat0], dim=1) 
        y0 = self.yolo0(h0)
        
        y0 = mkoutput(y0) 
        y1 = mkoutput(y1) 
        y2 = mkoutput(y2)

        if not self.training:  # inference
            B, C, H, W = x.shape 
            anch0 = torch.tensor([[10,13], [16,30], [33,23]], dtype=torch.float32, device=x.device).view(1, -1, 1, 1, 2)
            anch1 = torch.tensor([[30,61], [62,45], [59,119]], dtype=torch.float32, device=x.device).view(1, -1, 1, 1, 2)
            anch2 = torch.tensor([[116,90], [156,198], [373,326]], dtype=torch.float32, device=x.device).view(1, -1, 1, 1, 2)
            grid0 = self._make_grid(52, 52).to(x.device)
            grid1 = self._make_grid(26, 26).to(x.device)
            grid2 = self._make_grid(13, 13).to(x.device)

            y0[..., 0:2] = (y0[..., 0:2].sigmoid() + grid0) * 8  # xy
            y0[..., 2:4] = torch.exp(y0[..., 2:4]) * anch0 # wh
            y0[..., 4:] = y0[..., 4:].sigmoid()
            y0 = y0.view(B, -1, 85)

            y1[..., 0:2] = (y1[..., 0:2].sigmoid() + grid1) * 16  # xy
            y1[..., 2:4] = torch.exp(y1[..., 2:4]) * anch1 # wh
            y1[..., 4:] = y1[..., 4:].sigmoid()
            y1 = y1.view(B, -1, 85)

            y2[..., 0:2] = (y2[..., 0:2].sigmoid() + grid2) * 32  # xy
            y2[..., 2:4] = torch.exp(y2[..., 2:4]) * anch2 # wh
            y2[..., 4:] = y2[..., 4:].sigmoid()
            y2 = y2.view(B, -1, 85)

        return y0, y1, y2 

    @staticmethod
    def _make_grid(nx=20, ny=20):
        yv, xv = torch.meshgrid(torch.arange(ny), torch.arange(nx))
        return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
    @staticmethod
    def mkoutput(x, stride=8):
        B, C, H, W = x.shape 
        x = x.reshape(B, 3, 85, H, W).permute(0, 1, 3, 4, 2).contiguous()
        return x 
if __name__ == "__main__":
    model = YoloModel() 
    x = torch.randn([10, 3, 416, 416]) 
    y1, y2, y3 = model(x) 
    print(y1.shape, y2.shape, y3.shape)